import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from cbam_resnet import cbam_resnet50
from train import train_epoch, test

### 数据准备
 # CIFAR100的均值、标准差
CIFAR100_TRAIN_MEAN = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343)
CIFAR100_TRAIN_STD = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)

 # 数据增强
trainTransform = transforms.Compose([
    transforms.Resize((224, 224)), # 因为resnet50的输入是224×224
    transforms.RandomCrop(224, padding=4), # 随机裁剪
    transforms.RandomHorizontalFlip(), # 随机水平翻转
    transforms.RandomRotation(15), # 随机旋转
    transforms.ColorJitter(), # 图形属性随机更改 以上四步均是为了避免训练时过拟合
    transforms.ToTensor(),
    transforms.Normalize(CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD) # 标准化
])

validTransform = transforms.Compose([ 
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(CIFAR100_TRAIN_MEAN, CIFAR100_TRAIN_STD) # 标准化
])

 # CIFAR-100 dataset
train_dataset = torchvision.datasets.CIFAR100(root='./data/', train=True, transform=trainTransform, download=True)
test_dataset = torchvision.datasets.CIFAR100(root='./data/', train=False, transform=validTransform)
 
 # Data loader
batch_size=10
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

### 神经网络
#net=torchvision.models.resnet50(num_classes=100) # resnet50
net=cbam_resnet50(num_classes=100) #cbam_resnet50

### 损失函数和优化器
criterion = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = optim.Adam(net.parameters(), lr=1e-3)

### 数据记录
log_path = './logs/'
writer = SummaryWriter(log_path)

### 训练与测试
epoch_num = 10
lr0 = 1e-3
for epoch in range(epoch_num):
    current_lr = lr0 / 2**int(epoch/50)
    for param_group in optimizer.param_groups:
        param_group['lr'] = current_lr
    train_epoch(net, optimizer, train_loader, criterion, epoch, writer=writer)
    test(net, test_loader, criterion, epoch, writer=writer)
